#include <thrust/async/for_each.h>

#include <algorithm/analytical_model/algorithm/cuda/Utils.h>
#include <algorithm/analytical_model/algorithm/cuda/VectorMath.cuh>
#include <algorithm/analytical_model/algorithm/pipeline/AnalyticalPipeline.h>
#include <algorithm/shared/data/topology/geometry/Geometry.h>
#include <utils/Utils.h>

#include <cuda_runtime.h>
#include <cuda_runtime_api.h>

#include <bit>

using namespace solar;

auto solar::getFocusCenters(GeometryType receiver_type, const cuda::inter_op::Float3& receiver_size,
                            const cuda::inter_op::Float3& receiver_pos,
                            const cuda::inter_op::Float3& current_focus_center,
                            const cuda::inter_op::Float3* d_helio_pos,
                            const cuda::inter_op::Float3* d_focus_centers, std::size_t helio_count,
                            cudaStream_t stream) -> void
{
    cudaStream_t actual_stream = stream;
    if (stream == nullptr)
    {
        cudaStreamCreate(std::addressof(actual_stream));
    }
    if (d_focus_centers == nullptr)
    {
        CHECK_CUDA_ERRORS(cudaMallocAsync(std::addressof(d_focus_centers),
                                          sizeof(float3) * helio_count, actual_stream));
    }
    auto* d_focus_centers_copy = std::bit_cast<float3*>(d_focus_centers);
    auto* d_helio_pos_copy = std::bit_cast<float3*>(d_helio_pos);
    switch (receiver_type)
    {
    case GeometryType::kGeometryRectangle:
        [[fallthrough]];
    case GeometryType::kGeometryRectangle3D:
    {
        auto current_focus_center_copy = structCopy<float3>(current_focus_center);
        thrust::async::for_each(
            thrust::device.on(actual_stream), thrust::make_counting_iterator(std::size_t(0)),
            thrust::make_counting_iterator(helio_count),
            [current_focus_center_copy, d_focus_centers_copy] __device__(const std::size_t& index)
            { d_focus_centers_copy[index] = current_focus_center_copy; });
    }
    case GeometryType::kGeometryCylinder:
        [[fallthrough]];
    case GeometryType::kGeometryCylinder3D:
    {
        const auto receiver_size_copy = structCopy<float3>(receiver_size);
        const auto receiver_pos_copy = structCopy<float3>(receiver_pos);
        thrust::async::for_each(thrust::device.on(actual_stream),
                                thrust::make_counting_iterator(std::size_t(0)),
                                thrust::make_counting_iterator(helio_count),
                                [d_focus_centers_copy, d_helio_pos_copy, receiver_size_copy,
                                 receiver_pos_copy] __device__(const std::size_t& index)
                                {
                                    using namespace solar::cuda;
                                    const auto heliostat_position = d_helio_pos_copy[index];
                                    float3 dir = normalize(heliostat_position - receiver_pos_copy);
                                    float radius = receiver_size_copy.x / std::hypotf(dir.x, dir.z);
                                    float x = receiver_pos_copy.x + dir.x * radius;
                                    float z = receiver_pos_copy.z + dir.z * radius;
                                    d_focus_centers_copy[index] = {x, receiver_pos_copy.y, z};
                                });
    }
    }
    if (stream == nullptr)
    {
        CHECK_CUDA_ERRORS(cudaStreamSynchronize(actual_stream));
        CHECK_CUDA_ERRORS(cudaStreamDestroy(actual_stream));
    }
}
